package njust.csc.aecore.packs

import njust.csc.aecore.entity.APK
import njust.csc.aecore.utils.RegexUtils
import soot.*
import soot.Unit
import soot.jimple.AbstractJimpleValueSwitch
import soot.jimple.Stmt
import soot.jimple.VirtualInvokeExpr

/**
 * @Author ZhouHang
 * @Date 24/6/4 周二
 * 用于 Soot 提取 url，但因为速度慢已废弃
 */

interface Analyzer {
    fun addPack(apk: APK)
}

class UrlAnalyzer : Analyzer {
    override fun addPack(apk: APK) {
        PackManager.v().getPack("wjtp").remove("wjtp.url")
        val urls: MutableList<String> = mutableListOf()
        val transform = Transform("wjtp.url", object : SceneTransformer() {
            override fun internalTransform(p0: String?, p1: MutableMap<String, String>?) {
                Scene.v().classes.forEach{ sootClass ->
                    if (sootClass.name.startsWith(apk.packageName)) {
                        sootClass.methods.forEach{ sootMethod ->
                            try {
                                val body = sootMethod.retrieveActiveBody()
                                urls.addAll(RegexUtils.searchString(RegexUtils.urlMatch, body.toString()))
                            } catch (_: RuntimeException) {}
                        }
                    }
                }
                UrlAnalyzerResult.urls = urls
            }
        })

        PackManager.v().getPack("wjtp").add(transform)
    }
}

object UrlAnalyzerResult {
    var urls: MutableList<String> = mutableListOf()
}

class SdkAnalyzer : Analyzer {
    override fun addPack(apk: APK) {
        PackManager.v().getPack("wjtp").remove("wjtp.sdk")
        val transform = Transform("wjtp.sdk", object : SceneTransformer() {
            override fun internalTransform(p0: String?, p1: MutableMap<String, String>?) {
                Scene.v().classes.forEach{ sootClass ->
                    if (sootClass.name.startsWith(apk.packageName)) {
                        sootClass.methods.forEach{ sootMethod ->
                            try {
                                val body = sootMethod.retrieveActiveBody()
                                body.units.forEach { unit ->
                                    val sdk = findSdk(unit)
                                    if (sdk.isNotEmpty()) SdkAnalyzerResult.sdks.add(sdk)
                                }
                            } catch (_: RuntimeException) {}
                        }
                    }
                }
            }
        })
        PackManager.v().getPack("wjtp").add(transform)
    }

    private fun findSdk(unit: Unit): String {
        val stmt = unit as Stmt
        if (isVirtualInvokeStmt(stmt)) {
            val invokeExpr = stmt.invokeExpr as VirtualInvokeExpr
            val sdkResult = RegexUtils.sdkName.find(invokeExpr.toString())
            val sdkName = sdkResult?.value ?: ""
            return sdkName.drop(1).dropLast(1)
        }
        return ""
    }

    /**
     * 给定 stmt 是否为 Virtual Invoke 类型
     * @param stmt 给定 stmt
     */
    private fun isVirtualInvokeStmt(stmt: Stmt): Boolean {
        if (!stmt.containsInvokeExpr()) {
            return false
        }
        val flag = booleanArrayOf(false)
        val invokeExpr = stmt.invokeExpr
        invokeExpr.apply(object : AbstractJimpleValueSwitch<Any?>() {
            override fun caseVirtualInvokeExpr(v: VirtualInvokeExpr) {
                flag[0] = true
            }
        })
        return flag[0]
    }
}

object SdkAnalyzerResult {
    var sdks: MutableSet<String> = mutableSetOf()
}
